#!/usr/bin/env python3
"""
稳定版无人机导航系统模拟 V4.2
彻底重写MPC控制器，简化设计，提高稳定性
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import matplotlib.gridspec as gridspec
import math
import time
from queue import PriorityQueue
import cvxpy as cp
import warnings

# 忽略特定警告
warnings.filterwarnings("ignore", category=UserWarning, module="cvxpy")
warnings.filterwarnings("ignore", category=RuntimeWarning)

# 设置matplotlib使用支持中文的字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号

# 调试日志级别
DEBUG_LEVEL = 1  # 0: 无调试, 1: 基本调试, 2: 详细调试, 3: 非常详细


def debug_log(level, message):
    """调试日志函数"""
    if DEBUG_LEVEL >= level:
        print(f"[DEBUG{level}] {message}")


class PIDController:
    """PID控制器实现"""

    def __init__(self, kp, ki, kd, max_output=float('inf')):
        self.kp = kp  # 比例系数
        self.ki = ki  # 积分系数
        self.kd = kd  # 微分系数
        self.max_output = max_output  # 输出限制

        self.integral = np.zeros(3)
        self.previous_error = np.zeros(3)
        self.previous_time = time.time()
        self.error_history = []

    def compute(self, setpoint, current_value):
        """计算控制输出"""
        current_time = time.time()
        dt = current_time - self.previous_time
        if dt <= 0:
            return np.zeros(3)

        error = setpoint - current_value
        self.error_history.append(np.linalg.norm(error))

        # 比例项
        p_term = self.kp * error

        # 积分项
        self.integral += error * dt
        i_term = self.ki * self.integral

        # 微分项
        derivative = (error - self.previous_error) / dt
        d_term = self.kd * derivative

        # 计算总输出
        output = p_term + i_term + d_term

        # 限制输出
        output_norm = np.linalg.norm(output)
        if output_norm > 1e-6 and output_norm > self.max_output:
            output = output / output_norm * self.max_output

        # 更新状态
        self.previous_error = error
        self.previous_time = current_time

        debug_log(3, f"PID控制输出: {output}, 误差: {np.linalg.norm(error)}")

        return output

    def get_error_history(self):
        return self.error_history


class SimpleMPCController:
    """简化版模型预测控制器，专注于可行性"""

    def __init__(self, mass, horizon=5, dt=0.1):
        self.mass = mass  # 无人机质量
        self.horizon = horizon  # 预测步长
        self.dt = dt  # 时间步长

        # 简化的系统矩阵 (只考虑位置)
        self.A = np.eye(3)
        self.B = np.eye(3) * dt

        # 权重矩阵 - 调整权重
        self.Q = np.diag([10.0, 10.0, 15.0])  # 位置误差权重，Z方向权重更大
        self.R = np.diag([0.1, 0.1, 0.1])  # 控制输入权重

        # 控制输入约束
        self.u_min = -20.0
        self.u_max = 20.0

        # 状态约束
        self.x_min = np.array([-np.inf, -np.inf, 0.1])  # Z方向最小为0.1
        self.x_max = np.array([np.inf, np.inf, np.inf])

        # 优化问题变量
        self.x = cp.Variable((3, horizon + 1))
        self.u = cp.Variable((3, horizon))

        # 参考轨迹
        self.reference = None

        # 误差历史记录
        self.error_history = []

        # 备份PID控制器 - 增强参数
        self.backup_pid = PIDController(kp=np.array([4.0, 4.0, 6.0]),
                                        ki=np.array([0.05, 0.05, 0.1]),
                                        kd=np.array([2.0, 2.0, 3.0]),
                                        max_output=20.0)

        # 求解统计
        self.solve_count = 0
        self.fail_count = 0

    def set_reference(self, reference):
        """设置参考轨迹"""
        # 确保参考轨迹维度正确
        if reference.shape[0] != 3 or reference.shape[1] != self.horizon + 1:
            debug_log(1, f"参考轨迹维度错误: {reference.shape}，期望: (3, {self.horizon + 1})")
            # 尝试调整参考轨迹维度
            if reference.shape[0] != 3:
                reference = reference[:3, :]
            if reference.shape[1] < self.horizon + 1:
                # 填充参考轨迹
                last_point = reference[:, -1].reshape(3, 1)
                padding = np.repeat(last_point, self.horizon + 1 - reference.shape[1], axis=1)
                reference = np.hstack([reference, padding])
            else:
                reference = reference[:, :self.horizon + 1]

        self.reference = reference
        debug_log(2, f"设置参考轨迹: {reference[:, 0]}")

    def compute(self, current_state):
        """计算控制输入"""
        # 添加数值稳定性检查
        if np.any(np.isnan(current_state)) or np.any(np.isinf(current_state)):
            debug_log(1, "当前状态包含NaN或Inf值，使用备用PID")
            return self.backup_pid.compute(self.reference[:, 0] if self.reference is not None else np.zeros(3),
                                           current_state[0:3] if len(current_state) > 3 else current_state)

        if self.reference is None:
            debug_log(1, "参考轨迹未设置，返回零控制")
            return np.zeros(3)

        # 检查参考轨迹是否有效
        if np.any(np.isnan(self.reference)) or np.any(np.isinf(self.reference)):
            debug_log(1, "参考轨迹包含NaN或Inf值，使用备用PID")
            return self.backup_pid.compute(self.reference[:, 0],
                                           current_state[0:3] if len(current_state) > 3 else current_state)

        # 提取当前位置
        x0 = current_state[0:3] if len(current_state) > 3 else current_state

        # 构建成本函数和约束
        cost = 0
        constraints = [self.x[:, 0] == x0]

        for t in range(self.horizon):
            # 状态成本
            cost += cp.quad_form(self.x[:, t] - self.reference[:, t], self.Q)

            # 控制成本
            cost += cp.quad_form(self.u[:, t], self.R)

            # 系统动力学约束
            constraints += [self.x[:, t + 1] == self.A @ self.x[:, t] + self.B @ self.u[:, t]]

            # 控制输入约束
            constraints += [self.u[:, t] >= self.u_min, self.u[:, t] <= self.u_max]

            # 状态约束
            constraints += [self.x[:, t] >= self.x_min, self.x[:, t] <= self.x_max]

        # 终端成本
        cost += cp.quad_form(self.x[:, self.horizon] - self.reference[:, self.horizon], self.Q * 2)

        # 求解优化问题
        problem = cp.Problem(cp.Minimize(cost), constraints)

        self.solve_count += 1

        try:
            # 尝试使用ECOS求解器
            problem.solve(solver=cp.ECOS, verbose=False, max_iters=100, abstol=1e-3, reltol=1e-3)
        except Exception as e:
            try:
                # 如果ECOS失败，尝试OSQP
                problem.solve(solver=cp.OSQP, verbose=False, max_iter=4000)
            except Exception as e2:
                # 如果都失败，使用备用PID控制器
                self.fail_count += 1
                debug_log(1, f"所有求解器均失败: {e}, {e2}")
                debug_log(1,
                          f"MPC求解失败率: {self.fail_count}/{self.solve_count} = {self.fail_count / self.solve_count * 100:.2f}%")
                return self.backup_pid.compute(self.reference[:, 0], x0)

        if problem.status not in [cp.OPTIMAL, cp.OPTIMAL_INACCURATE]:
            self.fail_count += 1
            debug_log(1, f"MPC求解失败! 状态: {problem.status}")
            debug_log(1,
                      f"MPC求解失败率: {self.fail_count}/{self.solve_count} = {self.fail_count / self.solve_count * 100:.2f}%")
            return self.backup_pid.compute(self.reference[:, 0], x0)

        # 记录误差
        error = np.linalg.norm(x0 - self.reference[:, 0])
        self.error_history.append(error)

        # 返回第一个控制输入
        control_output = self.u[:, 0].value
        if control_output is None or not np.all(np.isfinite(control_output)):
            self.fail_count += 1
            debug_log(1, f"MPC控制输出无效: {control_output}")
            debug_log(1,
                      f"MPC求解失败率: {self.fail_count}/{self.solve_count} = {self.fail_count / self.solve_count * 100:.2f}%")
            return self.backup_pid.compute(self.reference[:, 0], x0)

        debug_log(2, f"MPC控制输出: {control_output}, 误差: {error}")
        debug_log(3, f"MPC求解成功率: {(self.solve_count - self.fail_count) / self.solve_count * 100:.2f}%")

        return control_output

    def get_error_history(self):
        return self.error_history


class SensorNoise:
    """传感器噪声模型"""

    def __init__(self):
        # GPS噪声参数 (位置测量)
        self.gps_position_std = 0.1  # 标准差 (米)
        self.gps_velocity_std = 0.05  # 标准差 (米/秒)

        # IMU噪声参数 (加速度测量)
        self.accel_bias = np.array([0.01, -0.02, 0.005])  # 加速度计偏置
        self.accel_std = 0.02  # 加速度计噪声标准差

        # 陀螺仪噪声参数 (角速度测量)
        self.gyro_bias = np.array([0.005, -0.003, 0.008])  # 陀螺仪偏置
        self.gyro_std = 0.01  # 陀螺仪噪声标准差

    def add_gps_noise(self, position, velocity):
        """添加GPS噪声"""
        noisy_position = position + np.random.normal(0, self.gps_position_std, 3)
        noisy_velocity = velocity + np.random.normal(0, self.gps_velocity_std, 3)
        return noisy_position, noisy_velocity

    def add_imu_noise(self, acceleration):
        """添加IMU噪声"""
        noisy_acceleration = acceleration + self.accel_bias + np.random.normal(0, self.accel_std, 3)
        return noisy_acceleration

    def add_gyro_noise(self, angular_velocity):
        """添加陀螺仪噪声"""
        noisy_angular_velocity = angular_velocity + self.gyro_bias + np.random.normal(0, self.gyro_std, 3)
        return noisy_angular_velocity


class AStarPlanner:
    """A*路径规划算法实现"""

    def __init__(self, grid_size=1.0):
        self.grid_size = grid_size

    def heuristic(self, a, b):
        """计算启发式函数（曼哈顿距离）"""
        return abs(a[0] - b[0]) + abs(a[1] - b[1]) + abs(a[2] - b[2])

    def plan(self, start, goal, obstacles):
        """执行A*路径规划"""
        # 将连续坐标转换为网格坐标
        start_grid = (int(start[0] / self.grid_size),
                      int(start[1] / self.grid_size),
                      int(start[2] / self.grid_size))
        goal_grid = (int(goal[0] / self.grid_size),
                     int(goal[1] / self.grid_size),
                     int(goal[2] / self.grid_size))

        # 创建障碍物网格
        obstacle_grid = set()
        for obs in obstacles:
            obs_grid = (int(obs[0] / self.grid_size),
                        int(obs[1] / self.grid_size),
                        int(obs[2] / self.grid_size))
            obstacle_grid.add(obs_grid)

        # A*算法
        frontier = PriorityQueue()
        frontier.put((0, start_grid))
        came_from = {}
        cost_so_far = {}
        came_from[start_grid] = None
        cost_so_far[start_grid] = 0

        # 可能的移动方向（26个方向）
        directions = []
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                for dz in [-1, 0, 1]:
                    if dx == 0 and dy == 0 and dz == 0:
                        continue
                    directions.append((dx, dy, dz))

        while not frontier.empty():
            current = frontier.get()[1]

            if current == goal_grid:
                break

            for dx, dy, dz in directions:
                next_node = (current[0] + dx, current[1] + dy, current[2] + dz)

                if next_node in obstacle_grid:
                    continue

                new_cost = cost_so_far[current] + math.sqrt(dx * dx + dy * dy + dz * dz)
                if next_node not in cost_so_far or new_cost < cost_so_far[next_node]:
                    cost_so_far[next_node] = new_cost
                    priority = new_cost + self.heuristic(goal_grid, next_node)
                    frontier.put((priority, next_node))
                    came_from[next_node] = current

        # 重建路径
        path = []
        current = goal_grid
        while current != start_grid:
            path.append(current)
            if current not in came_from:
                return []  # 没有找到路径
            current = came_from[current]
        path.append(start_grid)
        path.reverse()

        # 将网格坐标转换回连续坐标
        continuous_path = []
        for p in path:
            continuous_path.append(np.array([
                p[0] * self.grid_size + self.grid_size / 2,
                p[1] * self.grid_size + self.grid_size / 2,
                p[2] * self.grid_size + self.grid_size / 2
            ]))

        # 路径平滑处理
        if len(continuous_path) > 2:
            smoothed_path = self.smooth_path(continuous_path, obstacles)
            return smoothed_path
        return continuous_path

    def smooth_path(self, path, obstacles):
        """使用梯度下降法平滑路径"""
        smoothed = np.array(path)
        alpha = 0.1  # 平滑权重
        beta = 0.3  # 缩短权重

        for _ in range(50):  # 迭代次数
            for i in range(1, len(path) - 1):
                # 保持与障碍物的距离
                for obs in obstacles:
                    obs_dist = np.linalg.norm(smoothed[i] - obs)
                    if obs_dist < 2.0:  # 安全距离
                        smoothed[i] += 0.1 * (smoothed[i] - obs) / obs_dist

                # 平滑和缩短
                smoothed[i] += alpha * (path[i] - smoothed[i]) + beta * (
                        smoothed[i - 1] + smoothed[i + 1] - 2 * smoothed[i])

        return smoothed.tolist()


class Drone:
    def __init__(self, x=0, y=0, z=0, mass=1.0, controller_type="PID"):
        # 初始状态
        self.true_position = np.array([x, y, z], dtype=float)
        self.true_velocity = np.array([0, 0, 0], dtype=float)
        self.true_acceleration = np.array([0, 0, 0], dtype=float)

        # 传感器测量值 (带噪声)
        self.measured_position = np.array([x, y, z], dtype=float)
        self.measured_velocity = np.array([0, 0, 0], dtype=float)
        self.measured_acceleration = np.array([0, 0, 0], dtype=float)

        # 物理参数
        self.mass = mass
        self.max_thrust = 25.0  # 增加最大推力 (N)
        self.gravity = np.array([0, 0, -9.81])  # 重力加速度

        # 控制器
        self.pid_controller = PIDController(
            kp=np.array([4.0, 4.0, 6.0]),  # Z方向增益更高
            ki=np.array([0.05, 0.05, 0.1]),
            kd=np.array([2.0, 2.0, 3.0]),
            max_output=20.0
        )
        self.mpc_controller = SimpleMPCController(mass=mass, horizon=5, dt=0.1)
        self.controller_type = controller_type

        # 传感器噪声模型
        self.sensor_noise = SensorNoise()

        # 历史轨迹记录
        self.true_trajectory = []
        self.measured_trajectory = []

        # 参考轨迹 (用于MPC)
        self.reference_trajectory = None

        # 调试信息
        self.step_count = 0

    def set_controller(self, controller_type):
        """设置控制器类型"""
        self.controller_type = controller_type
        debug_log(1, f"设置控制器类型: {controller_type}")

    def set_reference_trajectory(self, reference):
        """设置参考轨迹 (用于MPC)"""
        if self.controller_type == "MPC":
            # 确保所有航点都是 NumPy 数组
            waypoints_np = [np.array(wp) for wp in reference]

            # 为每个航点段生成参考轨迹
            references = []
            for i in range(len(waypoints_np) - 1):
                start = waypoints_np[i]
                end = waypoints_np[i + 1]

                # 线性插值
                num_points = 5  # 减少插值点数量
                for j in range(num_points):
                    t = j / (num_points - 1)
                    point = start + t * (end - start)
                    references.append(point)

            debug_log(2, f"Generated {len(references)} reference points for MPC")

            # 转换为MPC控制器需要的格式 (3 x horizon+1)
            horizon = self.mpc_controller.horizon
            ref_states = np.zeros((3, horizon + 1))

            # 如果参考轨迹不够长，用最后一个点填充
            for i in range(horizon + 1):
                idx = min(i, len(references) - 1)
                ref_states[:, i] = references[idx]

            self.reference_trajectory = ref_states
            self.mpc_controller.set_reference(ref_states)

    def update_sensors(self):
        """更新传感器测量值 (添加噪声)"""
        # GPS测量 (位置和速度)
        self.measured_position, self.measured_velocity = self.sensor_noise.add_gps_noise(
            self.true_position, self.true_velocity
        )

        # IMU测量 (加速度)
        self.measured_acceleration = self.sensor_noise.add_imu_noise(
            self.true_acceleration
        )

        debug_log(3, f"传感器测量 - 真实位置: {self.true_position}, 测量位置: {self.measured_position}")

    def update(self, target_position, dt=0.1):
        """更新无人机状态"""
        self.step_count += 1

        # 更新传感器测量值
        self.update_sensors()

        # 根据控制器类型计算控制力
        if self.controller_type == "PID":
            # PID控制 - 使用测量值
            control_force = self.pid_controller.compute(target_position, self.measured_position)
            debug_log(2, f"使用PID控制器, 目标位置: {target_position}")
        else:
            # MPC控制 - 使用当前状态估计
            # 动态更新参考轨迹
            if self.reference_trajectory is not None:
                # 创建一个新的参考轨迹，从当前位置开始，平滑过渡到目标位置
                horizon = self.mpc_controller.horizon
                new_reference = np.zeros((3, horizon + 1))

                # 第一个点是当前位置
                new_reference[:, 0] = self.measured_position

                # 中间点线性插值
                for i in range(1, horizon):
                    t = i / horizon
                    new_reference[:, i] = self.measured_position + t * (target_position - self.measured_position)

                # 最后一个点是目标位置
                new_reference[:, horizon] = target_position

                self.mpc_controller.set_reference(new_reference)

            current_state = np.concatenate([self.measured_position, self.measured_velocity])
            control_force = self.mpc_controller.compute(current_state)
            debug_log(2, f"使用MPC控制器, 目标位置: {target_position}")

        # 计算总力 (控制力 + 重力)
        total_force = control_force + self.mass * self.gravity

        # 限制推力 - 修复除以零错误
        thrust_norm = np.linalg.norm(total_force)
        if thrust_norm > 1e-6 and thrust_norm > self.max_thrust:
            total_force = total_force / thrust_norm * self.max_thrust

        # 计算真实加速度 (F = ma)
        self.true_acceleration = total_force / self.mass

        # 更新真实速度和位置
        self.true_velocity += self.true_acceleration * dt
        self.true_position += self.true_velocity * dt

        # 确保无人机不会低于地面
        self.true_position[2] = max(0.1, self.true_position[2])

        # 记录轨迹
        self.true_trajectory.append(self.true_position.copy())
        self.measured_trajectory.append(self.measured_position.copy())

        # 改进的航点到达判断
        horizontal_distance = np.linalg.norm(target_position[0:2] - self.true_position[0:2])
        vertical_distance = abs(target_position[2] - self.true_position[2])

        # 水平距离小于0.5米且垂直距离小于0.3米视为到达
        reached = horizontal_distance < 0.5 and vertical_distance < 0.3

        debug_log(2, f"步数 {self.step_count}: 位置 {self.true_position}, 目标 {target_position}, 距离 {horizontal_distance}")

        return reached

    def get_true_position(self):
        return self.true_position.copy()

    def get_measured_position(self):
        return self.measured_position.copy()

    def get_true_trajectory(self):
        return np.array(self.true_trajectory)

    def get_measured_trajectory(self):
        return np.array(self.measured_trajectory)

    def get_controller_error_history(self):
        """获取控制器误差历史"""
        if self.controller_type == "PID":
            return self.pid_controller.get_error_history()
        else:
            return self.mpc_controller.get_error_history()


class NavigationSystem:
    def __init__(self, controller_type="PID"):
        # 创建障碍物
        self.obstacles = [
            np.array([5, 5, 2]),
            np.array([8, 3, 3]),
            np.array([12, 8, 4]),
            np.array([15, 12, 5]),
            np.array([7, 10, 3]),
            np.array([10, 15, 4]),
            np.array([13, 5, 3]),
            np.array([18, 8, 4])
        ]

        # 起点和终点
        self.start = np.array([0, 0, 1])
        self.goal = np.array([20, 15, 6])

        # 路径规划器
        self.planner = AStarPlanner(grid_size=1.0)

        # 规划路径
        self.waypoints = self.plan_path()

        # 当前目标航点索引
        self.current_waypoint_idx = 1  # 从第一个航点开始

        # 创建无人机实例
        self.drone = Drone(
            x=self.start[0],
            y=self.start[1],
            z=self.start[2],
            mass=1.0,  # 设置无人机质量
            controller_type=controller_type
        )

        # 为MPC设置参考轨迹
        if controller_type == "MPC":
            self.setup_mpc_reference()

        # 是否完成所有航点
        self.mission_complete = False

        # 调试信息
        self.update_count = 0

        # 打印规划路径
        debug_log(1, "规划路径:")
        for i, wp in enumerate(self.waypoints):
            debug_log(1, f"航点 {i}: {wp}")

    def plan_path(self):
        """使用A*算法规划路径"""
        debug_log(1, "Planning path...")
        path = self.planner.plan(self.start, self.goal, self.obstacles)

        if not path:
            debug_log(1, "No path found! Using straight line path")
            return [self.start, self.goal]

        debug_log(1, f"Path planning completed, total {len(path)} waypoints")
        for i, wp in enumerate(path):
            debug_log(2, f"Waypoint {i}: {wp}")

        return path

    def setup_mpc_reference(self):
        """为MPC设置参考轨迹"""
        # 确保所有航点都是 NumPy 数组
        waypoints_np = [np.array(wp) for wp in self.waypoints]

        # 为每个航点段生成参考轨迹
        references = []
        for i in range(len(waypoints_np) - 1):
            start = waypoints_np[i]
            end = waypoints_np[i + 1]

            # 线性插值
            num_points = 5  # 减少插值点数量
            for j in range(num_points):
                t = j / (num_points - 1)
                point = start + t * (end - start)
                references.append(point)

        debug_log(2, f"Generated {len(references)} reference points for MPC")
        self.drone.set_reference_trajectory(references)

    def update(self, dt=0.1):
        """更新导航系统状态"""
        self.update_count += 1

        if self.mission_complete:
            return True

        # 获取当前目标航点
        target = self.waypoints[self.current_waypoint_idx]

        debug_log(1, f"更新 {self.update_count}: 当前目标航点 {self.current_waypoint_idx}/{len(self.waypoints) - 1}: {target}")

        # 更新无人机位置
        reached = self.drone.update(target, dt)

        # 打印当前位置和规划路径
        current_pos = self.drone.get_true_position()
        debug_log(1, f"当前位置: {current_pos}, 规划路径点: {target}")

        # 如果到达当前航点，则转向下一个航点
        if reached:
            debug_log(1, f"到达航点 {self.current_waypoint_idx}: {target}")
            self.current_waypoint_idx += 1

            # 检查是否完成所有航点
            if self.current_waypoint_idx >= len(self.waypoints):
                debug_log(1, "任务完成!")
                self.mission_complete = True
                return True

        return False

    def get_drone_true_position(self):
        return self.drone.get_true_position()

    def get_drone_measured_position(self):
        return self.drone.get_measured_position()

    def get_waypoints(self):
        return self.waypoints

    def get_true_trajectory(self):
        return self.drone.get_true_trajectory()

    def get_measured_trajectory(self):
        return self.drone.get_measured_trajectory()

    def get_obstacles(self):
        return self.obstacles

    def get_controller_error_history(self):
        return self.drone.get_controller_error_history()


class Simulation:
    def __init__(self, controller_type="PID"):
        # 创建导航系统
        self.nav_system = NavigationSystem(controller_type)

        # 创建图形
        self.fig = plt.figure(figsize=(16, 12))

        # 使用GridSpec创建复杂的子图布局
        gs = gridspec.GridSpec(3, 2, figure=self.fig)

        # 3D轨迹图
        self.ax_3d = self.fig.add_subplot(gs[0:2, 0], projection='3d')

        # 误差图
        self.ax_error = self.fig.add_subplot(gs[0, 1])

        # 控制输入图
        self.ax_control = self.fig.add_subplot(gs[1, 1])

        # 传感器测量图
        self.ax_sensor = self.fig.add_subplot(gs[2, 0])

        # 性能指标图
        self.ax_performance = self.fig.add_subplot(gs[2, 1])

        # 设置3D图表标题和标签
        self.ax_3d.set_title(f'Drone Navigation Simulation - {controller_type} Controller')
        self.ax_3d.set_xlabel('X (m)')
        self.ax_3d.set_ylabel('Y (m)')
        self.ax_3d.set_zlabel('Z (m)')

        # 设置坐标轴范围
        self.ax_3d.set_xlim(-2, 25)
        self.ax_3d.set_ylim(-2, 20)
        self.ax_3d.set_zlim(0, 10)

        # 绘制起点和终点
        start = self.nav_system.start
        goal = self.nav_system.goal
        self.ax_3d.scatter(start[0], start[1], start[2], c='green', marker='o', s=100, label='Start')
        self.ax_3d.scatter(goal[0], goal[1], goal[2], c='red', marker='o', s=100, label='Goal')

        # 绘制航点
        waypoints = self.nav_system.get_waypoints()
        wp_x = [wp[0] for wp in waypoints]
        wp_y = [wp[1] for wp in waypoints]
        wp_z = [wp[2] for wp in waypoints]
        self.ax_3d.plot(wp_x, wp_y, wp_z, 'r--', alpha=0.5, label='Planned Path')

        # 绘制障碍物
        obstacles = self.nav_system.get_obstacles()
        for obs in obstacles:
            self.draw_cube(obs, 1.0, 'gray')

        # 初始化无人机位置标记
        self.true_marker, = self.ax_3d.plot([], [], [], 'bo', markersize=8, label='True Position')
        self.measured_marker, = self.ax_3d.plot([], [], [], 'rx', markersize=6, label='Measured Position')

        # 初始化轨迹线
        self.true_trajectory_line, = self.ax_3d.plot([], [], [], 'b-', alpha=0.7, linewidth=2, label='True Trajectory')
        self.measured_trajectory_line, = self.ax_3d.plot([], [], [], 'r-', alpha=0.5, linewidth=1,
                                                         label='Measured Trajectory')

        # 设置误差图
        self.ax_error.set_title('Controller Error')
        self.ax_error.set_xlabel('Time Step')
        self.ax_error.set_ylabel('Error (m)')
        self.ax_error.grid(True)
        self.error_line, = self.ax_error.plot([], [], 'b-')

        # 设置控制输入图
        self.ax_control.set_title('Control Input')
        self.ax_control.set_xlabel('Time Step')
        self.ax_control.set_ylabel('Control Force (N)')
        self.ax_control.grid(True)
        self.control_lines = [
            self.ax_control.plot([], [], label='Fx')[0],
            self.ax_control.plot([], [], label='Fy')[0],
            self.ax_control.plot([], [], label='Fz')[0]
        ]

        # 设置传感器测量图
        self.ax_sensor.set_title('Sensor Measurement vs True Value')
        self.ax_sensor.set_xlabel('Time Step')
        self.ax_sensor.set_ylabel('Position (m)')
        self.ax_sensor.grid(True)
        self.sensor_lines = [
            self.ax_sensor.plot([], [], label='True X')[0],
            self.ax_sensor.plot([], [], label='Measured X')[0],
            self.ax_sensor.plot([], [], label='True Y')[0],
            self.ax_sensor.plot([], [], label='Measured Y')[0]
        ]

        # 设置性能指标图
        self.ax_performance.set_title('Performance Metrics')
        self.ax_performance.set_xlabel('Metric')
        self.ax_performance.set_ylabel('Value')
        self.ax_performance.grid(True)

        # 添加图例
        self.ax_3d.legend()
        self.ax_control.legend()
        self.ax_sensor.legend()

        # 设置视角
        self.ax_3d.view_init(elev=20, azim=45)

        # 动画对象
        self.animation = None

        # 数据记录
        self.error_history = []
        self.control_history = []
        self.true_position_history = []
        self.measured_position_history = []
        self.time_steps = []

    def draw_cube(self, center, size, color):
        """绘制立方体障碍物"""
        # 计算立方体的边界
        r = size / 2
        x = [center[0] - r, center[0] + r]
        y = [center[1] - r, center[1] + r]
        z = [center[2] - r, center[2] + r]

        # 绘制立方体的框架
        from itertools import product
        vertices = list(product(x, y, z))

        # 定义立方体的边
        edges = [
            (0, 1), (0, 2), (0, 4),
            (1, 3), (1, 5),
            (2, 3), (2, 6),
            (3, 7),
            (4, 5), (4, 6),
            (5, 7),
            (6, 7)
        ]

        # 绘制每条边
        for edge in edges:
            start = vertices[edge[0]]
            end = vertices[edge[1]]
            self.ax_3d.plot3D(*zip(start, end), color=color, alpha=0.6)

    def update_plot(self, frame):
        """更新动画帧"""
        # 更新导航系统
        complete = self.nav_system.update(dt=0.1)

        # 获取无人机当前位置
        true_pos = self.nav_system.get_drone_true_position()
        measured_pos = self.nav_system.get_drone_measured_position()

        # 更新无人机位置标记
        self.true_marker.set_data([true_pos[0]], [true_pos[1]])
        self.true_marker.set_3d_properties([true_pos[2]])

        self.measured_marker.set_data([measured_pos[0]], [measured_pos[1]])
        self.measured_marker.set_3d_properties([measured_pos[2]])

        # 更新轨迹线
        true_trajectory = self.nav_system.get_true_trajectory()
        measured_trajectory = self.nav_system.get_measured_trajectory()

        if len(true_trajectory) > 1:
            self.true_trajectory_line.set_data(true_trajectory[:, 0], true_trajectory[:, 1])
            self.true_trajectory_line.set_3d_properties(true_trajectory[:, 2])

        if len(measured_trajectory) > 1:
            self.measured_trajectory_line.set_data(measured_trajectory[:, 0], measured_trajectory[:, 1])
            self.measured_trajectory_line.set_3d_properties(measured_trajectory[:, 2])

        # 更新误差图
        error_history = self.nav_system.get_controller_error_history()
        if error_history:
            self.error_line.set_data(range(len(error_history)), error_history)
            self.ax_error.relim()
            self.ax_error.autoscale_view()

        # 记录数据
        self.time_steps.append(len(self.time_steps))
        self.true_position_history.append(true_pos)
        self.measured_position_history.append(measured_pos)

        # 更新传感器测量图 (只显示最近50个点)
        recent_points = min(50, len(self.time_steps))
        if recent_points > 0:
            time_window = self.time_steps[-recent_points:]
            true_x = [p[0] for p in self.true_position_history[-recent_points:]]
            true_y = [p[1] for p in self.true_position_history[-recent_points:]]
            measured_x = [p[0] for p in self.measured_position_history[-recent_points:]]
            measured_y = [p[1] for p in self.measured_position_history[-recent_points:]]

            self.sensor_lines[0].set_data(time_window, true_x)
            self.sensor_lines[1].set_data(time_window, measured_x)
            self.sensor_lines[2].set_data(time_window, true_y)
            self.sensor_lines[3].set_data(time_window, measured_y)

            self.ax_sensor.relim()
            self.ax_sensor.autoscale_view()

        # 如果任务完成，停止动画并计算性能指标
        if complete:
            self.animation.event_source.stop()
            debug_log(1, "Simulation completed")
            self.calculate_performance()

        return (self.true_marker, self.measured_marker,
                self.true_trajectory_line, self.measured_trajectory_line,
                self.error_line, *self.sensor_lines)

    def calculate_performance(self):
        """计算性能指标"""
        true_trajectory = self.nav_system.get_true_trajectory()
        waypoints = self.nav_system.get_waypoints()

        if len(true_trajectory) < 2:
            return

        # 计算路径长度
        path_length = 0
        for i in range(1, len(true_trajectory)):
            path_length += np.linalg.norm(true_trajectory[i] - true_trajectory[i - 1])

        # 计算直线距离
        straight_distance = np.linalg.norm(waypoints[-1] - waypoints[0])

        # 计算效率 (直线距离/实际路径长度)
        efficiency = straight_distance / path_length if path_length > 0 else 0

        # 计算平均误差
        error_history = self.nav_system.get_controller_error_history()
        avg_error = np.mean(error_history) if error_history else 0

        # 计算最大误差
        max_error = np.max(error_history) if error_history else 0

        # 显示性能指标
        metrics = ['Path Length', 'Straight Distance', 'Efficiency', 'Avg Error', 'Max Error']
        values = [path_length, straight_distance, efficiency, avg_error, max_error]

        self.ax_performance.bar(metrics, values)
        self.ax_performance.tick_params(axis='x', rotation=45)

        # 添加数值标签
        for i, v in enumerate(values):
            self.ax_performance.text(i, v + 0.01, f'{v:.2f}', ha='center')

        # 打印性能指标
        debug_log(1, f"性能指标:")
        debug_log(1, f"  路径长度: {path_length:.2f} m")
        debug_log(1, f"  直线距离: {straight_distance:.2f} m")
        debug_log(1, f"  效率: {efficiency:.2f}")
        debug_log(1, f"  平均误差: {avg_error:.2f} m")
        debug_log(1, f"  最大误差: {max_error:.2f} m")

    def run(self):
        """运行模拟"""
        debug_log(1, "Starting drone navigation simulation...")
        debug_log(1, f"Controller type: {self.nav_system.drone.controller_type}")
        debug_log(1, f"Start point: {self.nav_system.start}")
        debug_log(1, f"Goal point: {self.nav_system.goal}")
        debug_log(1, f"Number of obstacles: {len(self.nav_system.get_obstacles())}")

        # 创建动画
        self.animation = FuncAnimation(
            self.fig, self.update_plot, frames=500,
            interval=50, blit=True, repeat=False
        )

        plt.tight_layout()
        plt.show()


# 比较不同控制器的性能
def compare_controllers():
    """比较PID和MPC控制器的性能"""
    debug_log(1, "Comparing performance of different controllers...")

    # 运行PID控制器模拟
    debug_log(1, "\n=== PID Controller ===")
    sim_pid = Simulation(controller_type="PID")
    sim_pid.run()

    # 运行MPC控制器模拟
    debug_log(1, "\n=== MPC Controller ===")
    sim_mpc = Simulation(controller_type="MPC")
    sim_mpc.run()


# 运行模拟
if __name__ == "__main__":
    # 先使用PID控制器测试基本功能
    debug_log(1, "First testing with PID controller to ensure basic functionality...")
    sim_pid = Simulation(controller_type="PID")
    sim_pid.run()

    # 然后尝试MPC控制器
    debug_log(1, "Now testing with MPC controller...")
    sim_mpc = Simulation(controller_type="MPC")
    sim_mpc.run()